import torch
import io
from typing import Optional, Dict, Any

class BasicEncryption:
    """基础加密模块"""
    def encrypt(self, data: bytes) -> bytes:
        return data  # 暂不实现实际加密

    def decrypt(self, data: bytes) -> bytes:
        return data

class ModelCompressor:
    """模型压缩工具"""
    def __init__(self, method: str = 'quantization'):
        self.method = method
        
    def compress(self, model_state: Dict[str, Any]) -> bytes:
        """多种压缩策略"""
        buffer = io.BytesIO()
        if self.method == 'quantization':
            quantized_state = {
                k: torch.quantize_per_tensor(v, 0.1, 0, torch.qint8) 
                for k, v in model_state.items()
            }
            torch.save(quantized_state, buffer)
        elif self.method == 'sparsity':
            sparse_state = {
                k: v.to_sparse() if v.dim() > 1 else v
                for k, v in model_state.items()
            }
            torch.save(sparse_state, buffer)
        else:
            torch.save(model_state, buffer)
        buffer.seek(0)
        return buffer.getvalue()

    def decompress(self, data: bytes) -> Dict[str, Any]:
        buffer = io.BytesIO(data)
        return torch.load(buffer)

class CommunicationProtocol:
    """通信协议主类"""
    def __init__(self, compression: bool = True, method: str = 'quantization'):
        self.encryption = BasicEncryption()
        self.compression = compression
        self.method = method
        self.compressor = ModelCompressor(method) if compression else None
        self.stats = {'sent': 0, 'received': 0, 'compressed': 0}

    def send(self, data: Dict[str, Any]) -> bytes:
        """发送数据"""
        if not isinstance(data, dict):
            raise ValueError("Data must be a dictionary")
            
        # 先压缩再加密
        if self.compression and self.compressor:
            compressed_data = self.compressor.compress(data)
            self.stats['compressed'] += 1
        else:
            buffer = io.BytesIO()
            torch.save(data, buffer)
            compressed_data = buffer.getvalue()
            
        encrypted = self.encryption.encrypt(compressed_data)
        self.stats['sent'] += 1
        return encrypted

    def receive(self, data: bytes) -> Dict[str, Any]:
        """接收数据"""
        if not isinstance(data, bytes):
            raise ValueError("Input must be bytes")
            
        decrypted = self.encryption.decrypt(data)
        
        if self.compression and self.compressor:
            model_state = self.compressor.decompress(decrypted)
        else:
            buffer = io.BytesIO(decrypted)
            model_state = torch.load(buffer)
            
        self.stats['received'] += 1
        return model_state
        
    def get_stats(self) -> Dict[str, int]:
        """获取通信统计信息"""
        return self.stats.copy()